import re
import string
import numpy as np
from typing import Optional, List, Union, Dict
from datetime import timedelta
# 修正：导入系统 Logger 类（而非直接导入 get_logger 函数）
from utils.logger import Logger
from utils.exceptions import ToolError
from middleware.redis_adapter import get_redis_adapter

# 修正：通过 Logger 类获取 logger 实例（指定 agent_name，task_id 用默认值）
logger = Logger.get_logger(agent_name="text_processing_tool")


class TextProcessingError(ToolError):
    """文本处理相关自定义异常"""
    pass


class TextProcessingTools:
    """系统级文本处理工具集，提供清洗、嵌入、格式处理等通用能力"""

    def __init__(self):
        """初始化文本处理配置"""
        # 中文标点符号集合（用于清洗）
        self.chinese_punctuation = "！？。，、；：""''（）《》【】￥…—·°‘’“”"
        # 嵌入缓存过期时间（24小时）
        self.embedding_cache_ttl = timedelta(hours=24)
        # 默认嵌入维度（可根据实际模型调整）
        self.default_embedding_dim = 768
        logger.info("TextProcessingTools initialized successfully")

    def text_cleaning(
            self,
            text: str,
            remove_punctuation: bool = True,
            remove_extra_spaces: bool = True,
            lower_case: bool = False,
            remove_special_chars: bool = True
    ) -> str:
        """
        通用文本清洗：支持多维度配置，适配中英文场景
        Args:
            text: 待清洗文本
            remove_punctuation: 是否移除标点符号（中英文）
            remove_extra_spaces: 是否移除多余空格（连续空格→单个空格，首尾空格）
            lower_case: 是否转为小写（仅对英文有效）
            remove_special_chars: 是否移除特殊字符（非文字、非数字、非标点）
        Returns:
            清洗后的文本
        """
        if not isinstance(text, str) or len(text.strip()) == 0:
            logger.warning("输入文本为空或非字符串类型，返回空字符串")
            return ""

        try:
            # 1. 基础处理：去除首尾空格、换行符
            cleaned = text.strip().replace("\n", " ").replace("\t", " ")

            # 2. 移除特殊字符（保留文字、数字、标点）
            if remove_special_chars:
                # 匹配中文、英文、数字、中英文标点
                pattern = r"[^\u4e00-\u9fa5a-zA-Z0-9\s" + re.escape(
                    string.punctuation + self.chinese_punctuation) + r"]"
                cleaned = re.sub(pattern, "", cleaned)

            # 3. 移除标点符号
            if remove_punctuation:
                # 合并中英文标点集合
                all_punctuation = string.punctuation + self.chinese_punctuation
                translator = str.maketrans("", "", all_punctuation)
                cleaned = cleaned.translate(translator)

            # 4. 移除多余空格（连续空格→单个空格）
            if remove_extra_spaces:
                cleaned = re.sub(r"\s+", " ", cleaned).strip()

            # 5. 转为小写（仅英文）
            if lower_case:
                cleaned = cleaned.lower()

            return cleaned
        except Exception as e:
            # 修正：日志输出适配系统格式（包含 agent_name 和默认 task_id）
            logger.error(f"文本清洗失败: text={text[:50]}..., error={str(e)}", exc_info=True)
            raise TextProcessingError(f"文本清洗失败: {str(e)}") from e

    def get_embedding(
            self,
            text: str,
            model_name: Optional[str] = None,
            cache: bool = True
    ) -> Optional[List[float]]:
        """
        生成文本嵌入（支持缓存优化，可扩展真实模型）
        Args:
            text: 待生成嵌入的文本
            model_name: 嵌入模型名称（用于区分不同模型的缓存）
            cache: 是否使用Redis缓存嵌入结果
        Returns:
            文本嵌入向量（列表形式）
        """
        # 预处理文本（清洗+去空）
        cleaned_text = self.text_cleaning(text)
        if not cleaned_text:
            logger.warning("清洗后的文本为空，无法生成嵌入")
            return None

        # 生成缓存键（包含模型名称，避免不同模型缓存冲突）
        cache_key = f"text_embedding:{model_name or 'default'}:{cleaned_text}"

        # 1. 尝试从缓存读取
        if cache:
            try:
                cached_embedding = redis_client.get(cache_key)
                if cached_embedding:
                    # 从字符串解析为列表（Redis存储格式："0.123,0.456,..."）
                    embedding = list(map(float, cached_embedding.decode("utf-8").split(",")))
                    logger.debug(f"从缓存加载嵌入: text={cleaned_text[:30]}...")
                    return embedding
            except Exception as e:
                logger.warning(f"Redis缓存读取失败，继续生成嵌入: {str(e)}")

        # 2. 生成嵌入（当前为基础实现，可替换为真实模型：Sentence-BERT、OpenAI API等）
        try:
            logger.debug(f"生成文本嵌入: text={cleaned_text[:30]}...")

            # 真实场景替换为：调用嵌入模型（如 sentence_transformers.encode 或 API）
            # 示例：使用随机向量作为占位（实际需替换为真实模型输出）
            embedding = np.random.randn(self.default_embedding_dim).tolist()
            # 归一化向量（符合嵌入向量标准）
            embedding = self._normalize_vector(embedding)

            # 3. 写入缓存（如果启用缓存）
            if cache:
                try:
                    # 向量转为字符串存储（"0.123,0.456,..."）
                    embedding_str = ",".join(map(str, embedding))
                    redis_client.set(
                        key=cache_key,
                        value=embedding_str,
                        ex=int(self.embedding_cache_ttl.total_seconds())
                    )
                    logger.debug(f"嵌入已写入缓存: cache_key={cache_key}")
                except Exception as e:
                    logger.warning(f"Redis缓存写入失败: {str(e)}")

            return embedding
        except Exception as e:
            logger.error(f"文本嵌入生成失败: text={cleaned_text[:50]}..., error={str(e)}", exc_info=True)
            raise TextProcessingError(f"文本嵌入生成失败: {str(e)}") from e

    def _normalize_vector(self, vector: List[float]) -> List[float]:
        """向量归一化（L2归一化，使向量长度为1）"""
        try:
            vector_np = np.array(vector)
            norm = np.linalg.norm(vector_np)
            if norm == 0:
                return vector
            return (vector_np / norm).tolist()
        except Exception as e:
            logger.error(f"向量归一化失败: error={str(e)}", exc_info=True)
            return vector

    def split_text_by_length(
            self,
            text: str,
            max_length: int = 512,
            overlap: int = 50,
            split_by_sentence: bool = True
    ) -> List[str]:
        """
        文本分段（适配大模型输入长度限制）
        Args:
            text: 待分段文本
            max_length: 每段最大长度（字符数）
            overlap: 段间重叠长度（避免语义断裂）
            split_by_sentence: 是否优先按句子分割（中文用。！？，；，英文用.?!;）
        Returns:
            分段后的文本列表
        """
        if not text or max_length <= 0:
            return []

        cleaned_text = self.text_cleaning(text, remove_punctuation=False, remove_extra_spaces=False)
        if len(cleaned_text) <= max_length:
            return [cleaned_text]

        segments = []
        start_idx = 0
        text_len = len(cleaned_text)

        if split_by_sentence:
            # 中文句子分隔符 + 英文句子分隔符
            sentence_separators = r"([。！？；，,.?!;])"
            # 按分隔符分割（保留分隔符）
            parts = re.split(sentence_separators, cleaned_text)
            # 重组句子（分隔符跟在句子后）
            sentences = []
            for i in range(0, len(parts), 2):
                sentence = parts[i]
                if i + 1 < len(parts):
                    sentence += parts[i + 1]
                if sentence.strip():
                    sentences.append(sentence)

            # 按句子累积分段
            current_segment = ""
            for sentence in sentences:
                if len(current_segment) + len(sentence) <= max_length:
                    current_segment += sentence
                else:
                    if current_segment:
                        segments.append(current_segment)
                        # 重叠处理：从当前段末尾截取overlap长度作为下一段开头
                        current_segment = current_segment[-overlap:] + sentence if overlap > 0 else sentence
                    else:
                        # 单个句子超长，直接截断
                        current_segment = sentence[:max_length]
            if current_segment:
                segments.append(current_segment)
        else:
            # 直接按字符长度分段（忽略句子结构）
            while start_idx < text_len:
                end_idx = start_idx + max_length
                segment = cleaned_text[start_idx:end_idx]
                segments.append(segment)
                # 下一段起始位置 = 当前结束位置 - 重叠长度
                start_idx = end_idx - overlap

        logger.debug(f"文本分段完成: 原文本长度={text_len}, 分段数={len(segments)}")
        return segments

    def count_text_stats(self, text: str) -> Dict[str, int]:
        """
        文本统计：字符数、字数（中文）、单词数（英文）、句子数
        Returns:
            统计结果字典
        """
        cleaned_text = self.text_cleaning(text, remove_punctuation=False, remove_extra_spaces=False)
        if not cleaned_text:
            logger.debug("输入文本为空，返回空统计结果")
            return {"char_count": 0, "chinese_word_count": 0, "english_word_count": 0, "sentence_count": 0}

        # 字符数（含标点、空格）
        char_count = len(cleaned_text)

        # 中文单词数（按字符计数，排除标点和空格）
        chinese_chars = re.findall(r"[\u4e00-\u9fa5]", cleaned_text)
        chinese_word_count = len(chinese_chars)

        # 英文单词数（按空格分割，排除纯标点）
        english_words = re.findall(r"[a-zA-Z]+", cleaned_text)
        english_word_count = len(english_words)

        # 句子数（按分隔符分割）
        sentence_separators = r"[。！？；,.?!;]"
        sentence_count = len(re.split(sentence_separators, cleaned_text.strip()))

        stats = {
            "char_count": char_count,
            "chinese_word_count": chinese_word_count,
            "english_word_count": english_word_count,
            "sentence_count": sentence_count
        }
        logger.debug(f"文本统计完成: {stats}")
        return stats


# 初始化 Redis 客户端（复用系统中间件）
try:
    redis_client = get_redis_adapter()  # 使用Redis适配器，支持集群模式
    logger.info("Redis客户端初始化成功（用于文本嵌入缓存）")
except Exception as e:
    logger.error(f"Redis客户端初始化失败: {str(e)}", exc_info=True)
    # 抛出异常终止工具初始化（缓存是可选功能，但Redis连接失败可能影响整体流程）
    raise TextProcessingError(f"Redis客户端初始化失败: {str(e)}") from e

# 单例工具实例（全局复用，避免重复初始化）
text_processing_tools = TextProcessingTools()

# 暴露核心函数（简化外部调用，无需通过实例）
text_cleaning = text_processing_tools.text_cleaning
get_embedding = text_processing_tools.get_embedding
split_text_by_length = text_processing_tools.split_text_by_length
count_text_stats = text_processing_tools.count_text_stats

# 测试入口（本地调试用）
if __name__ == "__main__":
    # 修正：测试时更新 logger 上下文（可选，方便区分测试任务）
    test_logger = Logger.update_context(task_id="text_processing_test", agent_name="text_processing_tool")
    test_logger.info("开始测试文本处理工具")

    test_text = """
    分布式多Agent系统（Distributed Multi-Agent System, DMAS）是由多个自主Agent组成的协作系统！
    每个Agent具有独立的决策能力，通过网络通信实现协同完成复杂任务。
    应用场景包括：智能调度、分布式计算、多机器人协作等。
    """

    # 测试文本清洗
    cleaned = text_cleaning(test_text)
    test_logger.info(f"清洗后文本: {cleaned}")

    # 测试嵌入生成
    embedding = get_embedding(test_text)
    test_logger.info(f"嵌入向量长度: {len(embedding) if embedding else None}")

    # 测试文本分段
    segments = split_text_by_length(test_text, max_length=100, overlap=20)
    test_logger.info(f"分段结果: {segments}")

    # 测试文本统计
    stats = count_text_stats(test_text)
    test_logger.info(f"文本统计: {stats}")

    test_logger.success("文本处理工具测试完成")